package com.coocaa.typhoon.api.v1.milvus.service;

import cn.hutool.core.lang.Assert;
import com.coocaa.typhoon.api.v1.milvus.model.HybridSearchCondition;
import io.milvus.param.dml.ranker.BaseRanker;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.AnnSearchReq;
import io.milvus.v2.service.vector.request.HybridSearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.request.data.SparseFloatVec;
import io.milvus.v2.service.vector.request.ranker.WeightedRanker;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * @author bijiahao
 * @date : 2025/9/25 10:17
 * @description
 */
@Slf4j
@Service
public class MilvusService {
    @Autowired
    private MilvusClientV2 client;

    public List<List<SearchResp.SearchResult>> hybridSearch(HybridSearchCondition condition) {
        Assert.notNull(condition, "HybridSearchCondition cannot be null");
        Assert.notEmpty(condition.getCollectionName(), "collectionName cannot be empty");
        if (null != condition.getQueryDense()) {
            Assert.notEmpty(condition.getVectorFieldName(), "vectorFieldName cannot be empty");
        }
        if (null != condition.getQuerySparse()) {
            Assert.notEmpty(condition.getSparseVectorFieldName(), "sparseVectorFieldName cannot be empty");
        }
        Assert.notEmpty(condition.getRerankerList(), "rerankerList cannot be empty");
        Assert.isTrue(condition.getDenseTopK() > 0, "denseTopK must be greater than 0");
        Assert.isTrue(condition.getSparseTopK() > 0, "sparseTopK must be greater than 0");
        Assert.isTrue(condition.getTopK() > 0, "topK must be greater than 0");
        Assert.notEmpty(condition.getDenseParams(), "denseParams cannot be empty");
        Assert.notEmpty(condition.getSparseParams(), "sparseParams cannot be empty");
        Assert.notEmpty(condition.getOutFields(), "outFields cannot be empty");

        List<AnnSearchReq> searchRequests = new ArrayList<>();
        boolean isCanRank = true;
        if (null != condition.getQueryDense()) {
            AnnSearchReq.AnnSearchReqBuilder denseBuilder = AnnSearchReq.builder()
                    .vectorFieldName(condition.getVectorFieldName())
                    .vectors(Collections.singletonList(new FloatVec(condition.getQueryDense())))
                    .params(condition.getDenseParams())
                    .topK(condition.getDenseTopK());
            if (StringUtils.isNotBlank(condition.getExpr())) {
                denseBuilder.expr(condition.getExpr());
            }
            searchRequests.add(denseBuilder.build());
        } else {
            isCanRank = false;
        }
        if (null != condition.getQuerySparse()) {
            AnnSearchReq.AnnSearchReqBuilder sparseBuilder = AnnSearchReq.builder()
                    .vectorFieldName(condition.getSparseVectorFieldName())
                    .vectors(Collections.singletonList(new SparseFloatVec(condition.getQuerySparse())))
                    .params(condition.getSparseParams())
                    .topK(condition.getSparseTopK());
            if (StringUtils.isNotBlank(condition.getExpr())) {
                sparseBuilder.expr(condition.getExpr());
            }
            searchRequests.add(sparseBuilder.build());
        } else {
            isCanRank = false;
        }
        HybridSearchReq.HybridSearchReqBuilder searchReqBuilder = HybridSearchReq.builder()
                .collectionName(condition.getCollectionName())
                .searchRequests(searchRequests)
                .outFields(condition.getOutFields())
                .topK(condition.getTopK());
//        if (isCanRank) {
//            BaseRanker reranker = new WeightedRanker(condition.getRerankerList());
//            searchReqBuilder.ranker(reranker);
//        }
//        BaseRanker reranker = new RRFRanker(100);
//        searchReqBuilder.ranker(reranker);
        WeightedRanker reranker = new WeightedRanker(Arrays.asList(0.7f));
            searchReqBuilder.ranker(reranker);
        HybridSearchReq hybridSearchReq = searchReqBuilder.build();
        SearchResp searchResp = client.hybridSearch(hybridSearchReq);
        return searchResp.getSearchResults();
    }

}
